import math
from torch.nn.functional import softmax
import torch
from torch import nn

# scaled dot product attention softmax(QK'/sqrt(d))V      additive attention softmax(W3tanh(W1K+W2Q))  * V
# muti-head
# self-attention


def sequence_mask(X, valid_len, value=0.0):
    """在序列中屏蔽不相关的项"""
    # bn, m  ;  bn

    # 1, m  <  vlen, 1  -> vlen, m   (vlen = bn)
    # 0, 1, 2, 3            len1, len1, len1, len1       true,true,true,false
    # 0, 1, 2, 3      <     len2, len2, len2, len2   ->  true,false,false,false
    # 0, 1, 2, 3            len3, len3, len3, len3       true,true,false,false
    # 将每行m后面的值置为0/-1e6
    maxlen = X.size(1)
    mask = torch.arange(maxlen, dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]
    X[~mask] = value
    return X


def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    # b,n,m
    if valid_lens is None:
        return softmax(X, dim=-1)
    else:
        shape = X.shape  # b, n, m
        if valid_lens.dim() == 1:  # b，重复n次，相同batch每个q的val_len相同
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        # X转为2维：bn,m
        X = sequence_mask(X.reshape(-1, shape[-1]), valid_lens, value=-1e6)
        return softmax(X.reshape(shape), dim=-1)


class DotProductAttention(nn.Module):
    """缩放点积注意力"""

    def __init__(self, dropout):
        super(DotProductAttention, self).__init__()
        self.dropout = nn.Dropout(dropout)

    # queries的形状：(batch_size，查询的个数n，d)
    # keys的形状：(batch_size，“键－值”对的个数m，d)
    # values的形状：(batch_size，“键－值”对的个数m，值的维度v)
    # valid_lens的形状:(batch_size)或者(batch_size，查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        # Q*trans(K)/sqrt(d)   b,n,d  x  b,d,m  ->  b,n,m
        scores = torch.bmm(queries, keys.transpose(1, 2)) / math.sqrt(d)
        # scores: b,n,m  ，将每行m后面的值置为-1e6，对应values每列后面的v无效，即后面几行的value无效
        self.attention_weights = masked_softmax(scores, valid_lens)
        # b,n,m x b,m,v -> b,n,v
        return torch.bmm(self.dropout(self.attention_weights), values)


def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # b,n,q*num -> b,n,num,q -> b,num,n,q -> b*num,n,q
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)
    X = X.permute(0, 2, 1, 3)
    return X.reshape(-1, X.shape[2], X.shape[3])


def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    # b*num,n,q -> b, num, n, q  ->  b, n, num*q  , 恢复原形状
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)


class MultiHeadAttention(nn.Module):
    """多头注意力"""
    # 输入维数均为num_hiddens
    # key_size/query_size/value_size 为每个头上的维数

    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.attention = DotProductAttention(dropout)
        self.W_q = nn.Linear(num_hiddens, query_size*num_heads, bias=bias)
        self.W_k = nn.Linear(num_hiddens, key_size*num_heads, bias=bias)
        self.W_v = nn.Linear(num_hiddens, value_size*num_heads, bias=bias)
        self.W_o = nn.Linear(value_size*num_heads, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # (batch_size，查询个数n 或者 “键－值”对的个数m，num_hiddens)
        # valid_lens: (batch_size)或(batch_size，查询的个数n)

        # 先线性层转换为多头需要的维数，再分割为多头处理
        # b,n,h -> b,n,q*num -> b*num,n,q
        # b,m,h -> b,m,k*num -> b*num,m,k
        # b,m,h -> b,m,v*num -> b*num,m,v     一般n=m, q=k

        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)

        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，
            # 然后如此复制第二项，然后诸如此类。
            # [[1,2,3],[5,5,5]] -> [[1,2,3],[1,2,3],[5,5,5],[5,5,5]]
            valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)

        # b*num,n,v
        output = self.attention(queries, keys, values, valid_lens)

        # b*num,n,v -> b,n,v*num
        output_concat = transpose_output(output, self.num_heads)
        # b,n,v*num -> b,n,h
        return self.W_o(output_concat)


class PositionalEncoding(nn.Module):
    """位置编码"""

    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的P
        self.P = torch.zeros((1, max_len, num_hiddens))
        # pos/10000^(2i/d)
        # [[0],         [0,2,4,...,h-1]
        #  [1],    /                     ->  len x h/2
        #  [2],
        #  [len]]
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        # 2i给sin，2i+1给cos
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        # b,l,h + 1,l,h
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)
